package tagging;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Vector;

public class Tag {
  public static final String START_TAG = "<s>";
  private static boolean ENABLE_ADD_ONE_TRANSITION = false;

  HashMap<String, Integer> tagCount_;
  HashMap<String, Integer> wordCount_;
  Vector<Integer> invWordCount_;     // Keeps how many words has appeared k times.

  int totCount_;
  
  public Tag() {
    tagCount_ = new HashMap<String, Integer>();
    wordCount_ = new HashMap<String, Integer>();
    invWordCount_ = new Vector<Integer>();
    for (int i = 0; i <= TrainingSet.UNKNOWN_THRESHOLD; ++i) {
      invWordCount_.add(0);
    }
    totCount_ = 0;
  }
  
  public void addWord(String word) {
    Integer count = wordCount_.get(word);
    
    if (count == null) {
      wordCount_.put(word, 1);
    } else {
      wordCount_.put(word, count + 1);
    }
    
    totCount_++;
  }
  
  public void addFollowingTag(String tag) {
    Integer tc = tagCount_.get(tag);
    
    if (tc == null) {
      tagCount_.put(tag, 1);
    } else {
      tagCount_.put(tag, tc + 1);
    }
  }
  
  public void process() {
    Iterator<Map.Entry<String, Integer>> it = null;

    it = wordCount_.entrySet().iterator();    
    while(it.hasNext()) {
      int count = it.next().getValue();
      if (count <= TrainingSet.UNKNOWN_THRESHOLD) {
        invWordCount_.set(count, invWordCount_.get(count) + 1);
      }
    }
  }
  
  public int getWordCount(String word) {
    Integer count = wordCount_.get(word);
    if (count == null) {
      return 0;
    }
    
    return count;
  }
  
  // Get Good-Turing count for a word which has count = c
  // It is assumed that c < UNKNOWN_THRESHOLD
  public double getGTWordCount(int c) {
    if (c == 0) {
      return (1.0) * (double)invWordCount_.get(1) / (double)(TrainingSet.getTotalWordCount() - totCount_);
    } else {
      if (invWordCount_.get(c+1) == 0) {
        return (c + 1.0);
      } else {
        return (c + 1.0) * (double)invWordCount_.get(c) / (double)invWordCount_.get(c+1);
      }
    }
  }
  
  public int getTagCount(String tag) {
    Integer tc = tagCount_.get(tag);
    if (tc == null) {
      tc = 0;
    }
    
    if (ENABLE_ADD_ONE_TRANSITION) tc++;
    
    return tc;
  }
  
  public int getTotalCount() {
    if (ENABLE_ADD_ONE_TRANSITION) return totCount_ + 46;  // number-of-tags = 46

    return totCount_;
  }
}
